package handlers

import (
	"fmt"
	"strings"
	"io/ioutil"
	"regexp"
	"bytes"
	x2j "github.com/basgys/goxml2json"
	"generalhandler/utils"
	"encoding/json"
	"generalhandler/conf"
	"time"
	"os"
	"path"
)

type GeneralHandler struct {

}

func (gh *GeneralHandler) Handle(bytesMsg []byte, log Logger, pub Publisher) error {
	start := time.Now().Unix()
	var mapMsg map[string]interface{}
	err := json.Unmarshal(bytesMsg, &mapMsg)
	if nil != err {
		log.Warn("failed to unmarshal %s, cuz %s", string(bytesMsg), err.Error())
		return nil
	}

	ifDomain, _ := mapMsg["domain"]
	strDomain, _ := ifDomain.(string)
	if 0 == len(strDomain) {
		log.Warn("found no domain in %s", string(bytesMsg))
		return nil
	}

	ifTkId, _ := mapMsg["tkid"]
	strTkId, _ := ifTkId.(string)
	if 0 == len(strTkId) {
		log.Warn("found no tkid in %s", string(bytesMsg))
		return nil
	}

	var ignoreMap map[string]interface{}
	ignoreMap = make(map[string]interface{})
	if ifIgnore, ok := mapMsg["ignore"]; ok{
		strIgnore, _ := ifIgnore.(string)
		utils.StringToMap(strIgnore,ignoreMap)
	}


	err = gh.handle(mapMsg, func(result map[string]interface{}, handlers []string) error {
		for k, v := range mapMsg {
			result[k] = v
		}
		result["start"] = start
		result["end"] = time.Now().Unix()
		bytesNewMsg, err := json.Marshal(result)
		if nil != err {
			return fmt.Errorf("failed to marshal result %v, cuz %s", result, err.Error())
		}
		for _, e := range handlers {
			if _, ok := ignoreMap[e]; ok{
				continue
			}
			pub.PublishMessage(e, bytesNewMsg)
		}
		return nil
	},pub,log)
	if nil != err {
		log.Error("failed to handle msg %s, cuz %s", string(bytesMsg), err.Error())
		return nil
	}

	return nil
}

type Callback func(result map[string]interface{}, handlers []string) error

func (gh *GeneralHandler) handle(msg map[string]interface{}, cb Callback,pub Publisher, log Logger) error {
	keys := strings.Split(conf.HandlerConfigInstance().InputKeys, "|")		
	vals := make([]string, len(keys))
	for i, e := range keys {
		val, ok := msg[e].(string)
		if !ok || 0 == len(val) {
			return fmt.Errorf("error in input [%s]", e)
		}
		vals[i] = val
	}

	vars := strings.Split(conf.HandlerConfigInstance().ReplaceVars, "|")
	if len(vars) != len(keys) {
		return fmt.Errorf("error, cuz len(vars) != len(keys)")
	}
	cas := conf.HandlerConfigInstance().CmdArgs
	for i, e := range vars {
		if 0 == len(e) {
			continue
		}
		cas = strings.Replace(cas, e, vals[i], -1)
	}
	cmdOutfile := conf.HandlerConfigInstance().CmdOutfile
	cas = strings.Replace(cas, "${OUTFILE}", cmdOutfile, -1)
	var args []string
	tmp := strings.Split(cas, "|")
	pat := regexp.MustCompile(`\$\{[^\{]+\}`)
	for _, e := range tmp {
		if pat.Match([]byte(e)) {
			continue
		}
		args = append(args, e)
	}
	cmdBin := conf.HandlerConfigInstance().CmdBin

	proxychainsBin := conf.HandlerConfigInstance().ProxychainsBin
	proxychainsConf := conf.HandlerConfigInstance().ProxychainsConf
	if 0 != len(proxychainsBin) && 0 != len(proxychainsConf) {
		confPath := path.Join(utils.GetSelfBinDirPath(), "proxychains.config")
		err := ioutil.WriteFile(confPath, []byte(strings.Replace(proxychainsConf, "|", "\n", -1)), 0644)
		if nil != err {
			return fmt.Errorf("failed to create proxychains.config, cuz %s", err.Error())
		}
		proxy_args := []string{"-f", confPath, cmdBin}
		cmdBin = proxychainsBin
		args = append(proxy_args, args...)
	}
	//fmt.Println(args)
	cmdLog := utils.NewCmdLogInstance(utils.GetSelfBinDirPath(), 10 * time.Minute)
	recorder := new(bytes.Buffer)
	err := utils.ExecCmd(conf.HandlerConfigInstance().CmdTimeout * 60, nil, recorder, cmdLog, cmdBin, args...)
	if nil != err {
		//recorder.WriteString(err.Error())
		//return fmt.Errorf("failed to exec %s, cuz %s", conf.HandlerConfigInstance().CmdBin, recorder.String())
		return fmt.Errorf("failed to exec %s", conf.HandlerConfigInstance().CmdBin)
	}

	out, err := ioutil.ReadFile(cmdOutfile)
	defer os.Remove(cmdOutfile)
	if nil != err {
		return fmt.Errorf("failed to read out file %s, cuz %s", cmdOutfile, err.Error())
	}
	// do nothing when empty result
	if 0 == len(out) {
		return nil
	}

	var mout map[string]interface{}
	mresult := make(map[string]interface{})
	var buf *bytes.Buffer
	if strings.HasSuffix(cmdOutfile, ".xml") {
		buf, err = x2j.Convert(bytes.NewReader(out))
		if nil != err {
			return fmt.Errorf("failed to x2j %s, cuz %s", string(out), err.Error())
		}
		if err := json.Unmarshal([]byte(strings.TrimSpace(buf.String())), &mout); nil != err {
			return fmt.Errorf("failed to json unmarshal %s, cuz %s", string(buf.Bytes()), err.Error())
		}
	} else if strings.HasSuffix(cmdOutfile, ".json") {
		buf = bytes.NewBuffer(out)
		if err := json.Unmarshal([]byte(strings.TrimSpace(buf.String())), &mout); nil != err {
			return fmt.Errorf("failed to json unmarshal %s, cuz %s", string(buf.Bytes()), err.Error())
		}
	} else {
		mout = map[string]interface{} {
			"raw_data": strings.TrimSpace(buf.String()),
		}
	}
	if error := recorder.String(); 0 != len(error) {
		mout["scan_error"] = error
	}
	mout["scan_type"] = conf.HandlerConfigInstance().ScanType
	mresult["scan_result"] = mout
	for _, handle := range conf.HandlerConfigInstance().OutputHandles {
		keys := strings.Split(handle.OutputKey, "|")
		if err = gh.travel(mresult, keys, handle.NewKey, handle.NextHandler, cb); nil != err {
			return fmt.Errorf("failed to handle output, cuz %s", err.Error())
		}
	}

	return nil
}



func (gh *GeneralHandler) travel(out interface{}, inputkeys []string, newkey string, nexts []string, cb Callback) error {
	if len(inputkeys) <= 0 {
		return nil // end
	}

	key := inputkeys[0]
	inputkeys = append(inputkeys[:0], inputkeys[1:]...)
	if 0 == len(inputkeys) {
		_map, ok := out.(map[string]interface{})
		if !ok {
			return fmt.Errorf("unexcept data %v", out)
		}
		ifval, ok := _map[key]
		if !ok {
			return nil // found no key
		}

		// rebuild message
		result := make(map[string]interface{})
		if 0 != len(newkey) {
			result[newkey] = ifval
		} else {
			result[key] = ifval
		}

		return cb(result, nexts) // handle
	}

	if strings.HasSuffix(key, "[]") {
		key = key[:len(key)-len("[]")]
		_map, ok := out.(map[string]interface{})
		if !ok {
			return fmt.Errorf("unexcept data %v", out)
		}
		ifval, ok := _map[key]
		if !ok {
			return nil // found no key
		}
		var _array []interface{}
		_array, ok = ifval.([]interface{})
		if ok {
		} else if item, ok := ifval.(map[string]interface{}); ok {
			_array = append(_array, item)
		} else {
			return fmt.Errorf("unexcept data %v", ifval)
		}
		for _, e := range _array {
			if err := gh.travel(e, inputkeys, newkey, nexts, cb); nil != err {
				return err
			}
		}
		return nil // empty array
	} else {
		_map, ok := out.(map[string]interface{})
		if !ok {
			return fmt.Errorf("unexcept data %v", out)
		}
		ifval, ok := _map[key]
		if ok {
			return gh.travel(ifval, inputkeys, newkey, nexts, cb)
		} else {
			return nil // found no key
		}
	}
}

